"""
    VerticalCovariance

Module implementing vertical correlation and covariance functions for GSI background error modeling.
This module handles vertical structure of atmospheric error covariances including:
- EOF (Empirical Orthogonal Function) decomposition
- Vertical correlation matrices
- Vertical mode transformations
- Level-dependent variance profiles

The vertical covariance structure is critical for properly representing how errors
correlate in the vertical direction, accounting for atmospheric dynamics such as:
- Stratification (stability layers)
- Vertical propagation of errors
- Scale-dependent vertical structures
"""
module VerticalCovariance

using LinearAlgebra
using Statistics

export VerticalCorrelationMatrix, EOFDecomposition, VerticalModeTransform
export compute_vertical_correlations, apply_vertical_mode_transform
export create_exponential_vertical_correlation, create_gaussian_vertical_correlation
export estimate_eof_from_ensemble, reconstruct_from_modes
export apply_vertical_correlation, apply_vertical_correlation_inverse
export validate_vertical_correlation, compute_mode_contributions

"""
    VerticalCorrelationMatrix{T<:AbstractFloat}

Vertical correlation matrix structure with analytical or empirical correlations.

# Fields
- `correlation_matrix::Matrix{T}`: Full vertical correlation matrix (nz × nz)
- `correlation_type::Symbol`: Type (:exponential, :gaussian, :empirical)
- `vertical_scale::T`: Characteristic vertical correlation length
- `nz::Int`: Number of vertical levels
- `eigenvalues::Vector{T}`: Eigenvalues of correlation matrix
- `eigenvectors::Matrix{T}`: Eigenvectors (for efficient application)
"""
struct VerticalCorrelationMatrix{T<:AbstractFloat}
    correlation_matrix::Matrix{T}
    correlation_type::Symbol
    vertical_scale::T
    nz::Int
    eigenvalues::Vector{T}
    eigenvectors::Matrix{T}

    function VerticalCorrelationMatrix{T}(corr_matrix::Matrix{T},
                                         corr_type::Symbol,
                                         vertical_scale::T) where T
        nz = size(corr_matrix, 1)

        # Ensure symmetry
        symmetric_corr = (corr_matrix + corr_matrix') / 2

        # Compute eigendecomposition for efficient application
        eigenvals, eigenvecs = eigen(symmetric_corr)

        # Ensure positive semi-definite (clip negative eigenvalues)
        eigenvals = max.(eigenvals, T(1e-12))

        new{T}(symmetric_corr, corr_type, vertical_scale, nz, eigenvals, eigenvecs)
    end
end

"""
    EOFDecomposition{T<:AbstractFloat}

Empirical Orthogonal Function (EOF) decomposition for vertical structure.

EOFs provide an optimal basis for representing vertical variability using
fewer modes than the full vertical resolution.

# Fields
- `eof_modes::Matrix{T}`: EOF basis vectors (nz × n_modes)
- `eigenvalues::Vector{T}`: Variance explained by each mode
- `explained_variance_ratio::Vector{T}`: Fraction of variance per mode
- `n_modes::Int`: Number of retained modes
- `nz::Int`: Number of vertical levels
- `mode_truncation_threshold::T`: Threshold for mode retention
"""
struct EOFDecomposition{T<:AbstractFloat}
    eof_modes::Matrix{T}
    eigenvalues::Vector{T}
    explained_variance_ratio::Vector{T}
    n_modes::Int
    nz::Int
    mode_truncation_threshold::T

    function EOFDecomposition{T}(eof_modes::Matrix{T},
                                eigenvals::Vector{T};
                                truncation_threshold::T = T(0.99)) where T
        nz, n_modes = size(eof_modes)

        # Compute explained variance ratio
        total_variance = sum(eigenvals)
        explained_ratio = eigenvals ./ total_variance

        # Cumulative variance for mode selection
        cumulative_variance = cumsum(explained_ratio)

        # Determine number of modes to retain
        n_retained = findfirst(cumulative_variance .>= truncation_threshold)
        n_retained = isnothing(n_retained) ? n_modes : n_retained

        new{T}(eof_modes, eigenvals, explained_ratio, n_retained, nz, truncation_threshold)
    end
end

"""
    VerticalModeTransform{T<:AbstractFloat}

Complete vertical mode transformation system combining correlations and EOF.

# Fields
- `eof_decomposition::EOFDecomposition{T}`: EOF basis
- `correlation_matrix::VerticalCorrelationMatrix{T}`: Vertical correlations
- `mode_variance::Vector{T}`: Variance in each mode
- `transform_matrix::Matrix{T}`: Combined transform (correlation × EOF)
"""
struct VerticalModeTransform{T<:AbstractFloat}
    eof_decomposition::EOFDecomposition{T}
    correlation_matrix::VerticalCorrelationMatrix{T}
    mode_variance::Vector{T}
    transform_matrix::Matrix{T}

    function VerticalModeTransform{T}(eof_decomp::EOFDecomposition{T},
                                     vcorr::VerticalCorrelationMatrix{T}) where T
        # Compute mode variances (diagonal of E^T * C * E)
        eof_matrix = eof_decomp.eof_modes[:, 1:eof_decomp.n_modes]
        mode_variance = vec(sum((vcorr.correlation_matrix * eof_matrix) .* eof_matrix, dims=1))

        # Combined transform matrix
        transform_matrix = vcorr.correlation_matrix * eof_matrix

        new{T}(eof_decomp, vcorr, mode_variance, transform_matrix)
    end
end

"""
    create_exponential_vertical_correlation(nz::Int, vertical_scale::T;
                                           level_spacing::Vector{T} = ones(T, nz)) where T

Create exponential vertical correlation matrix.

Correlation structure: C(Δz) = exp(-|Δz|/L_vert)

# Arguments
- `nz::Int`: Number of vertical levels
- `vertical_scale::T`: Vertical correlation length (in level units or pressure)
- `level_spacing::Vector{T}`: Spacing between levels (default uniform)

# Returns
- `VerticalCorrelationMatrix{T}`: Exponential correlation structure
"""
function create_exponential_vertical_correlation(nz::Int, vertical_scale::T;
                                                 level_spacing::Vector{T} = ones(T, nz)) where T
    corr_matrix = Matrix{T}(undef, nz, nz)

    # Compute cumulative distances between levels
    cumulative_distance = zeros(T, nz)
    for k in 2:nz
        cumulative_distance[k] = cumulative_distance[k-1] + level_spacing[k-1]
    end

    # Fill correlation matrix
    for i in 1:nz, j in 1:nz
        distance = abs(cumulative_distance[i] - cumulative_distance[j])
        corr_matrix[i,j] = exp(-distance / vertical_scale)
    end

    return VerticalCorrelationMatrix{T}(corr_matrix, :exponential, vertical_scale)
end

"""
    create_gaussian_vertical_correlation(nz::Int, vertical_scale::T;
                                        level_spacing::Vector{T} = ones(T, nz)) where T

Create Gaussian vertical correlation matrix.

Correlation structure: C(Δz) = exp(-(Δz)²/(2L²))

# Arguments
- `nz::Int`: Number of vertical levels
- `vertical_scale::T`: Vertical correlation length
- `level_spacing::Vector{T}`: Spacing between levels

# Returns
- `VerticalCorrelationMatrix{T}`: Gaussian correlation structure
"""
function create_gaussian_vertical_correlation(nz::Int, vertical_scale::T;
                                             level_spacing::Vector{T} = ones(T, nz)) where T
    corr_matrix = Matrix{T}(undef, nz, nz)

    # Compute cumulative distances
    cumulative_distance = zeros(T, nz)
    for k in 2:nz
        cumulative_distance[k] = cumulative_distance[k-1] + level_spacing[k-1]
    end

    # Fill correlation matrix with Gaussian structure
    for i in 1:nz, j in 1:nz
        distance = abs(cumulative_distance[i] - cumulative_distance[j])
        corr_matrix[i,j] = exp(-(distance^2) / (2 * vertical_scale^2))
    end

    return VerticalCorrelationMatrix{T}(corr_matrix, :gaussian, vertical_scale)
end

"""
    compute_vertical_correlations(vertical_profiles::Matrix{T};
                                  correlation_type::Symbol = :empirical) where T

Compute vertical correlations from a set of vertical profiles (e.g., from ensemble).

# Arguments
- `vertical_profiles::Matrix{T}`: Profiles matrix (nz × n_samples)
- `correlation_type::Symbol`: Method (:empirical, :exponential_fit, :gaussian_fit)

# Returns
- `VerticalCorrelationMatrix{T}`: Estimated vertical correlation structure
"""
function compute_vertical_correlations(vertical_profiles::Matrix{T};
                                      correlation_type::Symbol = :empirical) where T
    nz, n_samples = size(vertical_profiles)

    if correlation_type == :empirical
        # Compute sample correlation matrix
        # Center the profiles
        centered_profiles = vertical_profiles .- mean(vertical_profiles, dims=2)

        # Compute covariance
        cov_matrix = (centered_profiles * centered_profiles') / (n_samples - 1)

        # Convert to correlation
        std_devs = sqrt.(diag(cov_matrix))
        corr_matrix = cov_matrix ./ (std_devs * std_devs')

        # Estimate characteristic scale from decay
        vertical_scale = estimate_vertical_scale(corr_matrix)

        return VerticalCorrelationMatrix{T}(corr_matrix, :empirical, vertical_scale)

    elseif correlation_type == :exponential_fit
        # Fit exponential correlation to empirical data
        empirical_corr = compute_vertical_correlations(vertical_profiles, correlation_type=:empirical)

        # Fit exponential to lag-1 correlation
        lag1_corr = mean([empirical_corr.correlation_matrix[i,i+1] for i in 1:nz-1])
        fitted_scale = -one(T) / log(lag1_corr)

        return create_exponential_vertical_correlation(nz, fitted_scale)

    elseif correlation_type == :gaussian_fit
        # Fit Gaussian correlation
        empirical_corr = compute_vertical_correlations(vertical_profiles, correlation_type=:empirical)

        # Fit Gaussian to lag-1 correlation
        lag1_corr = mean([empirical_corr.correlation_matrix[i,i+1] for i in 1:nz-1])
        fitted_scale = sqrt(-one(T) / (2 * log(lag1_corr)))

        return create_gaussian_vertical_correlation(nz, fitted_scale)

    else
        throw(ArgumentError("Unknown correlation type: $correlation_type"))
    end
end

"""
    estimate_vertical_scale(corr_matrix::Matrix{T}) where T

Estimate characteristic vertical correlation length from correlation matrix.

# Method
Finds the lag at which correlation drops to 1/e ≈ 0.368

# Returns
- `T`: Estimated vertical correlation scale
"""
function estimate_vertical_scale(corr_matrix::Matrix{T}) where T
    nz = size(corr_matrix, 1)

    # Average correlation as function of lag
    avg_corr = zeros(T, nz)
    for lag in 0:nz-1
        correlations = [corr_matrix[i, min(i+lag, nz)] for i in 1:nz-lag]
        avg_corr[lag+1] = mean(correlations)
    end

    # Find lag where correlation drops to 1/e
    target_corr = T(1/ℯ)
    scale_idx = findfirst(avg_corr .< target_corr)

    if isnothing(scale_idx)
        # If never drops below 1/e, use decay rate
        return T(nz / 3)  # Default to ~1/3 of domain
    else
        return T(scale_idx - 1)
    end
end

"""
    estimate_eof_from_ensemble(ensemble_profiles::Matrix{T};
                               truncation_threshold::T = T(0.99)) where T

Estimate EOF decomposition from ensemble vertical profiles.

# Arguments
- `ensemble_profiles::Matrix{T}`: Ensemble profiles (nz × n_ensemble)
- `truncation_threshold::T`: Cumulative variance threshold for mode retention

# Returns
- `EOFDecomposition{T}`: EOF basis with selected modes
"""
function estimate_eof_from_ensemble(ensemble_profiles::Matrix{T};
                                   truncation_threshold::T = T(0.99)) where T
    nz, n_ens = size(ensemble_profiles)

    # Center ensemble
    ensemble_mean = mean(ensemble_profiles, dims=2)
    centered_ensemble = ensemble_profiles .- ensemble_mean

    # Compute covariance matrix
    cov_matrix = (centered_ensemble * centered_ensemble') / (n_ens - 1)

    # Eigendecomposition (EOFs are eigenvectors)
    eigenvals, eigenvecs = eigen(cov_matrix)

    # Sort by decreasing eigenvalue (variance)
    sorted_indices = sortperm(eigenvals, rev=true)
    sorted_eigenvals = eigenvals[sorted_indices]
    sorted_eigenvecs = eigenvecs[:, sorted_indices]

    # Normalize EOFs
    for i in 1:nz
        sorted_eigenvecs[:, i] ./= norm(sorted_eigenvecs[:, i])
    end

    return EOFDecomposition{T}(sorted_eigenvecs, sorted_eigenvals,
                              truncation_threshold=truncation_threshold)
end

"""
    apply_vertical_mode_transform(transform::VerticalModeTransform{T},
                                  column::Vector{T};
                                  inverse::Bool = false) where T

Apply vertical mode transformation to a vertical column.

# Arguments
- `transform::VerticalModeTransform{T}`: Transformation system
- `column::Vector{T}`: Vertical profile (nz) or mode coefficients (n_modes)
- `inverse::Bool`: If true, transform from mode space to physical space

# Returns
- `Vector{T}`: Transformed profile

# Transformations
- Forward: x_mode = E^T * C * x_physical
- Inverse: x_physical = C^(-1) * E * x_mode
"""
function apply_vertical_mode_transform(transform::VerticalModeTransform{T},
                                      column::Vector{T};
                                      inverse::Bool = false) where T
    nz = transform.correlation_matrix.nz
    n_modes = transform.eof_decomposition.n_modes

    if inverse
        # Mode space → Physical space
        if length(column) != n_modes
            throw(DimensionMismatch("Column length $(length(column)) != n_modes $n_modes"))
        end

        # Apply: x = E * x_mode (simplified, without correlation inverse)
        eof_matrix = transform.eof_decomposition.eof_modes[:, 1:n_modes]
        return eof_matrix * column

    else
        # Physical space → Mode space
        if length(column) != nz
            throw(DimensionMismatch("Column length $(length(column)) != nz $nz"))
        end

        # Apply: x_mode = E^T * x
        eof_matrix = transform.eof_decomposition.eof_modes[:, 1:n_modes]
        return eof_matrix' * column
    end
end

"""
    apply_vertical_correlation(vcorr::VerticalCorrelationMatrix{T},
                              column::Vector{T}) where T

Apply vertical correlation matrix to a vertical profile.

# Arguments
- `vcorr::VerticalCorrelationMatrix{T}`: Vertical correlation structure
- `column::Vector{T}`: Input vertical profile (nz)

# Returns
- `Vector{T}`: Correlated vertical profile
"""
function apply_vertical_correlation(vcorr::VerticalCorrelationMatrix{T},
                                   column::Vector{T}) where T
    if length(column) != vcorr.nz
        throw(DimensionMismatch("Column length $(length(column)) != nz $(vcorr.nz)"))
    end

    # Efficient application using eigendecomposition
    # C * x = V * Λ * V^T * x
    transformed = vcorr.eigenvectors' * column
    scaled = vcorr.eigenvalues .* transformed
    return vcorr.eigenvectors * scaled
end

"""
    apply_vertical_correlation_inverse(vcorr::VerticalCorrelationMatrix{T},
                                       column::Vector{T};
                                       regularization::T = T(1e-8)) where T

Apply inverse vertical correlation matrix: C^(-1) * x

# Arguments
- `vcorr::VerticalCorrelationMatrix{T}`: Vertical correlation structure
- `column::Vector{T}`: Input vertical profile
- `regularization::T`: Regularization for numerical stability

# Returns
- `Vector{T}`: Result of C^(-1) * x
"""
function apply_vertical_correlation_inverse(vcorr::VerticalCorrelationMatrix{T},
                                           column::Vector{T};
                                           regularization::T = T(1e-8)) where T
    if length(column) != vcorr.nz
        throw(DimensionMismatch("Column length $(length(column)) != nz $(vcorr.nz)"))
    end

    # C^(-1) = V * Λ^(-1) * V^T
    transformed = vcorr.eigenvectors' * column
    # Regularized inverse of eigenvalues
    inv_eigenvals = one(T) ./ (vcorr.eigenvalues .+ regularization)
    scaled = inv_eigenvals .* transformed
    return vcorr.eigenvectors * scaled
end

"""
    reconstruct_from_modes(eof_decomp::EOFDecomposition{T},
                          mode_coefficients::Vector{T}) where T

Reconstruct vertical profile from mode coefficients.

# Arguments
- `eof_decomp::EOFDecomposition{T}`: EOF decomposition
- `mode_coefficients::Vector{T}`: Coefficients for each mode

# Returns
- `Vector{T}`: Reconstructed vertical profile (nz)
"""
function reconstruct_from_modes(eof_decomp::EOFDecomposition{T},
                               mode_coefficients::Vector{T}) where T
    n_coeff = length(mode_coefficients)
    n_modes_available = size(eof_decomp.eof_modes, 2)

    if n_coeff > n_modes_available
        throw(ArgumentError("More coefficients ($n_coeff) than available modes ($n_modes_available)"))
    end

    # Reconstruction: x = E * coeffs
    eof_matrix = eof_decomp.eof_modes[:, 1:n_coeff]
    return eof_matrix * mode_coefficients
end

"""
    compute_mode_contributions(eof_decomp::EOFDecomposition{T},
                              profile::Vector{T}) where T

Compute contribution of each mode to a vertical profile.

# Returns
- `Vector{T}`: Mode coefficients showing contribution of each EOF
"""
function compute_mode_contributions(eof_decomp::EOFDecomposition{T},
                                   profile::Vector{T}) where T
    if length(profile) != eof_decomp.nz
        throw(DimensionMismatch("Profile length $(length(profile)) != nz $(eof_decomp.nz)"))
    end

    # Project onto EOFs: coeffs = E^T * x
    n_modes = min(eof_decomp.n_modes, size(eof_decomp.eof_modes, 2))
    eof_matrix = eof_decomp.eof_modes[:, 1:n_modes]
    return eof_matrix' * profile
end

"""
    validate_vertical_correlation(vcorr::VerticalCorrelationMatrix{T}) where T

Validate vertical correlation matrix properties.

# Checks
- Symmetry
- Positive semi-definiteness
- Diagonal elements = 1
- Physical correlation decay

# Returns
- `Dict{Symbol,Bool}`: Validation results
"""
function validate_vertical_correlation(vcorr::VerticalCorrelationMatrix{T}) where T
    results = Dict{Symbol,Bool}()

    # Check symmetry
    results[:symmetric] = issymmetric(vcorr.correlation_matrix)

    # Check positive semi-definite (all eigenvalues >= 0)
    results[:positive_semidefinite] = all(vcorr.eigenvalues .>= -1e-10)

    # Check diagonal elements ≈ 1
    diag_elements = diag(vcorr.correlation_matrix)
    results[:diagonal_unit] = all(abs.(diag_elements .- 1) .< 1e-6)

    # Check correlation decay (off-diagonal elements < 1)
    nz = vcorr.nz
    off_diag_max = maximum([abs(vcorr.correlation_matrix[i,j])
                           for i in 1:nz, j in 1:nz if i != j])
    results[:proper_decay] = off_diag_max < 1.0

    results[:all_valid] = all(values(results))

    return results
end

end # module VerticalCovariance
